"""
WebSocket路由
处理实时通信
"""

import json
import asyncio
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Optional

from ..services.notification_service import notification_service
from ..logger import logger

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    client_id: Optional[str] = Query(None, description="客户端ID")
):
    """WebSocket连接端点"""
    await websocket.accept()
    
    # 添加连接
    await notification_service.add_websocket_connection(websocket, client_id)
    
    try:
        # 发送欢迎消息
        welcome_message = {
            "type": "welcome",
            "message": "Connected to Model Subscription System",
            "client_id": client_id,
            "timestamp": __import__('time').time()
        }
        await websocket.send_text(json.dumps(welcome_message))
        
        # 保持连接并处理消息
        while True:
            try:
                # 接收客户端消息
                data = await websocket.receive_text()
                message = json.loads(data)
                
                # 处理客户端消息
                await handle_client_message(websocket, client_id, message)
                
            except WebSocketDisconnect:
                break
            except json.JSONDecodeError:
                # 发送错误消息
                error_message = {
                    "type": "error",
                    "message": "Invalid JSON format",
                    "timestamp": __import__('time').time()
                }
                await websocket.send_text(json.dumps(error_message))
            except Exception as e:
                logger.error(f"Error handling WebSocket message: {e}")
                error_message = {
                    "type": "error",
                    "message": "Internal server error",
                    "timestamp": __import__('time').time()
                }
                await websocket.send_text(json.dumps(error_message))
                
    except WebSocketDisconnect:
        pass
    except Exception as e:
        logger.error(f"WebSocket connection error: {e}")
    finally:
        # 移除连接
        await notification_service.remove_websocket_connection(websocket, client_id)
        logger.debug(f"WebSocket connection closed for client {client_id}")


async def handle_client_message(websocket: WebSocket, client_id: Optional[str], message: dict):
    """处理客户端消息"""
    try:
        message_type = message.get('type', '')
        
        if message_type == 'ping':
            # 心跳检测
            pong_message = {
                "type": "pong",
                "timestamp": __import__('time').time()
            }
            await websocket.send_text(json.dumps(pong_message))
            
        elif message_type == 'subscribe':
            # 订阅特定事件
            await handle_subscribe_message(websocket, client_id, message)
            
        elif message_type == 'unsubscribe':
            # 取消订阅
            await handle_unsubscribe_message(websocket, client_id, message)
            
        elif message_type == 'get_status':
            # 获取状态信息
            await handle_status_request(websocket, client_id, message)
            
        else:
            # 未知消息类型
            error_message = {
                "type": "error",
                "message": f"Unknown message type: {message_type}",
                "timestamp": __import__('time').time()
            }
            await websocket.send_text(json.dumps(error_message))
            
    except Exception as e:
        logger.error(f"Error handling client message: {e}")


async def handle_subscribe_message(websocket: WebSocket, client_id: Optional[str], message: dict):
    """处理订阅消息"""
    try:
        # 这里可以实现更细粒度的订阅控制
        # 目前简单返回确认消息
        
        response = {
            "type": "subscribe_ack",
            "message": "Subscription confirmed",
            "data": message.get('data', {}),
            "timestamp": __import__('time').time()
        }
        await websocket.send_text(json.dumps(response))
        
    except Exception as e:
        logger.error(f"Error handling subscribe message: {e}")


async def handle_unsubscribe_message(websocket: WebSocket, client_id: Optional[str], message: dict):
    """处理取消订阅消息"""
    try:
        response = {
            "type": "unsubscribe_ack",
            "message": "Unsubscription confirmed",
            "data": message.get('data', {}),
            "timestamp": __import__('time').time()
        }
        await websocket.send_text(json.dumps(response))
        
    except Exception as e:
        logger.error(f"Error handling unsubscribe message: {e}")


async def handle_status_request(websocket: WebSocket, client_id: Optional[str], message: dict):
    """处理状态请求"""
    try:
        # 获取连接统计信息
        stats = await notification_service.get_connection_stats()
        
        # 获取客户端的订阅信息
        from ..storage import data_manager
        client_subscriptions = []
        
        if client_id:
            subscriptions = await data_manager.get_subscriptions()
            client_subscriptions = [
                {
                    "id": sub.id,
                    "model_name": sub.model_name,
                    "server_id": sub.server_id,
                    "status": sub.status
                }
                for sub in subscriptions if sub.client_id == client_id
            ]
        
        response = {
            "type": "status",
            "data": {
                "client_id": client_id,
                "connection_stats": stats,
                "subscriptions": client_subscriptions,
                "server_time": __import__('time').time()
            },
            "timestamp": __import__('time').time()
        }
        await websocket.send_text(json.dumps(response))
        
    except Exception as e:
        logger.error(f"Error handling status request: {e}")


@router.websocket("/ws/admin")
async def admin_websocket_endpoint(websocket: WebSocket):
    """管理员WebSocket连接端点"""
    await websocket.accept()
    
    # 添加管理员连接
    await notification_service.add_websocket_connection(websocket, "admin")
    
    try:
        # 发送管理员欢迎消息
        welcome_message = {
            "type": "admin_welcome",
            "message": "Connected as administrator",
            "timestamp": __import__('time').time()
        }
        await websocket.send_text(json.dumps(welcome_message))
        
        # 定期发送系统状态
        status_task = asyncio.create_task(send_periodic_status(websocket))
        
        # 处理管理员消息
        while True:
            try:
                data = await websocket.receive_text()
                message = json.loads(data)
                await handle_admin_message(websocket, message)
                
            except WebSocketDisconnect:
                break
            except Exception as e:
                logger.error(f"Error in admin WebSocket: {e}")
                
    except WebSocketDisconnect:
        pass
    finally:
        # 取消状态任务
        if 'status_task' in locals():
            status_task.cancel()
            
        # 移除连接
        await notification_service.remove_websocket_connection(websocket, "admin")
        logger.debug("Admin WebSocket connection closed")


async def send_periodic_status(websocket: WebSocket):
    """定期发送系统状态"""
    try:
        while True:
            await asyncio.sleep(30)  # 每30秒发送一次状态
            
            # 获取系统状态
            stats = await notification_service.get_connection_stats()
            
            from ..storage import data_manager
            models = await data_manager.get_models()
            subscriptions = await data_manager.get_subscriptions()
            servers = await data_manager.get_servers()
            
            status_message = {
                "type": "system_status",
                "data": {
                    "models_count": len(models),
                    "subscriptions_count": len(subscriptions),
                    "servers_count": len(servers),
                    "active_subscriptions": len([s for s in subscriptions if s.status == "active"]),
                    "online_servers": len([s for s in servers if s.status == "online"]),
                    "connection_stats": stats
                },
                "timestamp": __import__('time').time()
            }
            
            await websocket.send_text(json.dumps(status_message))
            
    except asyncio.CancelledError:
        pass
    except Exception as e:
        logger.error(f"Error sending periodic status: {e}")


async def handle_admin_message(websocket: WebSocket, message: dict):
    """处理管理员消息"""
    try:
        message_type = message.get('type', '')
        
        if message_type == 'broadcast':
            # 广播消息到所有客户端
            broadcast_data = message.get('data', {})
            await notification_service.send_websocket_message(broadcast_data)
            
            response = {
                "type": "broadcast_ack",
                "message": "Broadcast sent",
                "timestamp": __import__('time').time()
            }
            await websocket.send_text(json.dumps(response))
            
        elif message_type == 'get_connections':
            # 获取连接信息
            stats = await notification_service.get_connection_stats()
            response = {
                "type": "connections_info",
                "data": stats,
                "timestamp": __import__('time').time()
            }
            await websocket.send_text(json.dumps(response))
            
        else:
            # 处理其他管理员命令
            response = {
                "type": "admin_response",
                "message": f"Command {message_type} received",
                "timestamp": __import__('time').time()
            }
            await websocket.send_text(json.dumps(response))
            
    except Exception as e:
        logger.error(f"Error handling admin message: {e}")
